package org.junyee.demo.security.ec;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.reflect.Field;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.Signature;
import java.security.interfaces.ECKey;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.text.MessageFormat;

import org.junyee.demo.security.ec.PemUtils.PemEncodedKeySpec;
import org.junyee.demo.security.exception.KeyException;
import org.springframework.util.Assert;
import org.springframework.util.Base64Utils;
import org.springframework.util.ReflectionUtils;

import lombok.Getter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import sun.security.x509.AlgorithmId;
import sun.security.util.DerOutputStream;

import sun.security.util.DerValue;
import sun.security.util.ECUtil;

/**
 * link https://lapo.it/asn1js ASN.1 DER PKCS8 Format
 */
@Getter
@ToString
@Slf4j
@SuppressWarnings("restriction")
public class EcPem {
    public static final String SECURITY_SUN_EC_PARAM_OPTIONAL_DISABLE = "security.sun.ec.params.optional.disable";
    private static final String FORMAT_PUBLIC_PEM = "-----BEGIN PUBLIC KEY-----\r\n{0}\r\n-----END PUBLIC KEY-----\r\n";
    private static final String FORMAT_PRIVATE_PEM = "-----BEGIN PRIVATE KEY-----\r\n{0}\r\n-----END PRIVATE KEY-----\r\n";

    private ECPrivateKey privateKey;
    private ECPublicKey publicKey;
    private String publicPem;
    private String privatePem;

    private EcPem() {

    }

    private int getSHALengthNum() {
        ECKey ecKey = privateKey == null ? publicKey : privateKey;
        int num = ((ecKey.getParams().getOrder().bitLength() + 7) / 8 * 8);
        num = num <= 128 ? 1 : num > 512 ? 512 : num;
        return num;
    }

    public Signature generateSignature() {
        try {
            String algname = "SHA" + getSHALengthNum() + "withECDSA";
            if (log.isDebugEnabled()) {
                log.debug("The algorithm is [{}] from PEM.", algname);
            }
            return Signature.getInstance(algname);
        } catch (NoSuchAlgorithmException e) {
            throw new KeyException(e);
        }
    }

    public static EcPem buildFromPrivateKeyPem(String pem) {
        if (!PemUtils.isPkcs8PrivateKeyFormat(pem)) {
            throw new KeyException("Not PKCS8 Private Key PEM");
        }
        return buildFromPem(pem);
    }

    public static EcPem buildFromPem(String pem) {
        try {
            PemEncodedKeySpec keySpec = PemUtils.getKeySpec(pem);
            PKCS8EncodedKeySpec privateKey2 = keySpec.getPrivateKey();
            ECPrivateKey ecPrivateKey = null;
            ECPublicKey ecPublicKey = null;
            KeyFactory keyFactory = EcUtils.KEY_FACTORY;
            if (privateKey2 != null) {
                ecPrivateKey = (ECPrivateKey) keyFactory.generatePrivate(privateKey2);
            }
            X509EncodedKeySpec publicKey2 = keySpec.getPublicKey();
            if (publicKey2 != null) {
                ecPublicKey = (ECPublicKey) keyFactory.generatePublic(publicKey2);
            }
            return build(ecPrivateKey, ecPublicKey);
        } catch (Exception e) {
            throw new KeyException(e);
        }
    }

    public static EcPem build(ECPrivateKey privateKey, ECPublicKey publicKey) {
        Assert.isTrue(privateKey != null || publicKey != null, "key requried");
        EcPem ecPem = new EcPem();
        ecPem.privateKey = privateKey;
        ecPem.publicKey = publicKey;
        if (publicKey != null) {
            ecPem.publicPem = ecPem.exportPkcs8PublickeyPem();
        }
        if (privateKey != null) {
            ecPem.privatePem = ecPem.exportPkcs8PrivatekeyPem();
        }
        return ecPem;
    }

    private String exportPkcs8PublickeyPem() {
        String pemSource = Base64Utils.encodeToString(publicKey.getEncoded());
        return MessageFormat.format(FORMAT_PUBLIC_PEM, PemUtils.enhance(pemSource));
    }

    private String exportPkcs8PrivatekeyPem() {
        DerOutputStream out = new DerOutputStream();
        try {
            out.write(DerValue.tag_Sequence, buidPrivateKeyAsn1());
            byte[] bytes = out.toByteArray();
            String pemSource = Base64Utils.encodeToString(bytes);
            return MessageFormat.format(FORMAT_PRIVATE_PEM, PemUtils.enhance(pemSource));
        } catch (IOException e) {
            throw new KeyException(e);
        }
    }

    private DerOutputStream buidPrivateKeyAsn1() {
        DerOutputStream privateKeyAsn1 = new DerOutputStream();
        try {
            // 0: DER Version
            privateKeyAsn1.putInteger(0);
            DerOutputStream algDerStream = buidAlgorithmIdAsn1();
            byte[] algorithmIdBytes = algDerStream.toByteArray();
            privateKeyAsn1.write(algorithmIdBytes);
            privateKeyAsn1.write(DerValue.tag_OctetString, buidSec1(algorithmIdBytes));
            return privateKeyAsn1;
        } catch (IOException e) {
            throw new KeyException(e);
        }
    }

    private DerOutputStream buidAlgorithmIdAsn1() {
        DerOutputStream derOutputStream = new DerOutputStream();
        try {
            getAlgorithmId().encode(derOutputStream);
            return derOutputStream;
        } catch (IOException e) {
            throw new KeyException(e);
        }
    }

    private boolean enableEcParamOptional() {
        return Boolean.parseBoolean(System.getProperty(SECURITY_SUN_EC_PARAM_OPTIONAL_DISABLE, "true"));
    }

    private AlgorithmId getAlgorithmId() {
        Field field = ReflectionUtils.findField(privateKey.getClass(), "algid");
        ReflectionUtils.makeAccessible(field);
        AlgorithmId algid = (AlgorithmId) ReflectionUtils.getField(field, privateKey);
        return algid;
    }

    private byte[] getAlgorithmIdDataBytes(byte[] algorithmIdBytes) {
        try {
            DerValue derValue = new DerValue(new ByteArrayInputStream(algorithmIdBytes));
            if (derValue.tag != DerValue.tag_Sequence) {
                throw new KeyException("invalid algorithmId bytes");
            }
            return derValue.getDataBytes();
        } catch (IOException e) {
            throw new KeyException(e);
        }
    }

    private byte[] buidSec1(byte[] algorithmIdBytes) {
        try {
            byte[] sArr = privateKey.getS().toByteArray();
            int numOctets = (privateKey.getParams().getOrder().bitLength() + 7) / 8;
            byte[] sOctets = new byte[numOctets];
            int inPos = Math.max(sArr.length - sOctets.length, 0);
            int outPos = Math.max(sOctets.length - sArr.length, 0);
            int length = Math.min(sArr.length, sOctets.length);
            System.arraycopy(sArr, inPos, sOctets, outPos, length);
            DerOutputStream out = new DerOutputStream();
            out.putInteger(1); // version 1
            out.putOctetString(sOctets);
            if (!enableEcParamOptional()) {
                // 0xA0: DER [0] ec params Optional
                out.putDerValue(new DerValue((byte) 0xA0, getAlgorithmIdDataBytes(algorithmIdBytes)));
            }
            // 0xA1: DER [1] public key Optional
            out.putDerValue(new DerValue((byte) 0xA1, buildPublicKeyOptional()));
            DerOutputStream seq = new DerOutputStream();
            seq.write(DerValue.tag_Sequence, out);
            return seq.toByteArray();
        } catch (IOException exc) {
            throw new KeyException(exc);
        }
    }

    private byte[] buildPublicKeyOptional() {
        byte[] key = ECUtil.encodePoint(publicKey.getW(), publicKey.getParams().getCurve());
        DerOutputStream out = new DerOutputStream();
        try {
            out.putBitString(key);
        } catch (IOException e) {
            throw new KeyException(e);
        }
        return out.toByteArray();
    }
}
